/**
 * Copyright (c) 2024 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the
 * "License"). Please refer to the License for details. You may not use this
 * file except in compliance with the License. THIS SOFTWARE IS PROVIDED ON AN
 * "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS
 * FOR A PARTICULAR PURPOSE. See LICENSE in the root of the software repository
 * for the full text of the License.
 */

#ifndef EXAMPLES_SORT_TOPK_CUSTOM_H
#define EXAMPLES_SORT_TOPK_CUSTOM_H
#include "kernel_operator.h"


namespace MyCustomKernel {
struct VecTiling {
    TopkTiling topKTilingData;
    uint32_t k;
    uint32_t outter;
    uint32_t inner;
    uint32_t n;
    uint32_t minsize;
    bool isLargest;
};

constexpr uint8_t K_FLOAT = 8;
constexpr uint8_t K_HALF = 16;
constexpr uint8_t LOCAL_BYTES = 32;

template <typename T, typename U, bool isInitIndex = false,
        bool isHasfinish = true, bool isReuseSrc = false,
        bool topkMode = false, bool tmpLocal = true>
class KernelTopK {
public:
    __aicore__ inline KernelTopK() {}
    __aicore__ inline void Init(GM_ADDR srcGmValue, GM_ADDR srcGmIndex,
                            GM_ADDR finishGm, GM_ADDR dstGmValue,
                            GM_ADDR dstGmIndex, VecTiling tilingData)
    {
        outter = tilingData.outter;
        inner = tilingData.inner;
        n = tilingData.n;
        isSmallMode = topkMode;
        tmplocalBytes = tilingData.minsize;
        topKTilingData = tilingData.topKTilingData;
        k = tilingData.k;
        // 计算k_pad
        if (sizeof(T) == sizeof(float)) {
            k_pad = (k + K_FLOAT - 1) / K_FLOAT * K_FLOAT;
        } else {
            k_pad = (k + K_HALF - 1) / K_HALF * K_HALF;
        }
        kpad_index = (k + K_FLOAT) / K_FLOAT * K_FLOAT;
        isLargest = tilingData.isLargest;
        inDataSize = inner * outter;
        outValueDataSize = k_pad * outter;
        outIndexDataSize = kpad_index * outter;

        inputdexDataSize = inner;
        if (topkMode == true) {
            inputdexDataSize = inDataSize;
        }

        finishLocalBytes = outter * sizeof(bool);
        if (finishLocalBytes % LOCAL_BYTES != 0) {
            finishLocalBytes = (finishLocalBytes + LOCAL_BYTES - 1) / LOCAL_BYTES * LOCAL_BYTES;
        }

        srcGlobal1.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(srcGmValue), inDataSize);
        srcGlobal2.SetGlobalBuffer(reinterpret_cast<__gm__ U *>(srcGmIndex), inputdexDataSize);
        srcGlobal3.SetGlobalBuffer(reinterpret_cast<__gm__ bool *>(finishGm), finishLocalBytes / sizeof(bool));
        dstGlobal1.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(dstGmValue), outValueDataSize);
        dstGlobal2.SetGlobalBuffer(reinterpret_cast<__gm__ U *>(dstGmIndex), outIndexDataSize);

        pipe.InitBuffer(inQueueX1, 1, inDataSize * sizeof(T));
        pipe.InitBuffer(inQueueX2, 1, inputdexDataSize * sizeof(U));
        pipe.InitBuffer(inQueueX3, 1, finishLocalBytes);
        pipe.InitBuffer(outQueueY1, 1, outValueDataSize * sizeof(T));
        pipe.InitBuffer(outQueueY2, 1, outIndexDataSize * sizeof(U));
    }
    __aicore__ inline void Process() {
        CopyIn();
        Compute();
        CopyOut();
    }

private:
    __aicore__ inline void CopyIn() {
        AscendC::LocalTensor<T> srcLocalValue = inQueueX1.AllocTensor<T>();
        AscendC::LocalTensor<U> srcLocalIndex = inQueueX2.AllocTensor<U>();
        AscendC::LocalTensor<bool> srcLocalFinish = inQueueX3.AllocTensor<bool>();
        AscendC::DataCopy(srcLocalValue, srcGlobal1, inDataSize);
        AscendC::DataCopy(srcLocalIndex, srcGlobal2, inputdexDataSize);
        AscendC::DataCopy(srcLocalFinish, srcGlobal3, finishLocalBytes / sizeof(bool));

        inQueueX1.EnQue(srcLocalValue);
        inQueueX2.EnQue(srcLocalIndex);
        inQueueX3.EnQue(srcLocalFinish);
    }
    __aicore__ inline void Compute() {
        AscendC::LocalTensor<T> dstLocalValue = outQueueY1.AllocTensor<T>();
        AscendC::LocalTensor<U> dstLocalIndex = outQueueY2.AllocTensor<U>();

        AscendC::LocalTensor<T> srcLocalValue = inQueueX1.DeQue<T>();
        AscendC::LocalTensor<U> srcLocalIndex = inQueueX2.DeQue<U>();
        AscendC::LocalTensor<bool> srcLocalFinish = inQueueX3.DeQue<bool>();

        auto newTopkMode = AscendC::TopKMode::TOPK_NORMAL;
        if (isSmallMode) {
            newTopkMode = AscendC::TopKMode::TOPK_NSMALL;
        }
        auto topKInfo = AscendC::TopKInfo();
        topKInfo.outter = outter;
        topKInfo.inner = inner;
        topKInfo.n = n;

        T scalar1(0);
        AscendC::Duplicate<T>(dstLocalValue, scalar1, outValueDataSize);
        U scalar2(0);
        AscendC::Duplicate<U>(dstLocalIndex, scalar2, outIndexDataSize);
        if (!tmpLocal) {
            if (isSmallMode) {
                AscendC::TopK<T, isInitIndex, isHasfinish, isReuseSrc, AscendC::TopKMode::TOPK_NSMALL>(dstLocalValue,
                    dstLocalIndex, srcLocalValue, srcLocalIndex, srcLocalFinish, k, topKTilingData, topKInfo,
                    isLargest);
            } else {
                AscendC::TopK<T, isInitIndex, isHasfinish, isReuseSrc, AscendC::TopKMode::TOPK_NORMAL>(dstLocalValue,
                    dstLocalIndex, srcLocalValue, srcLocalIndex, srcLocalFinish, k, topKTilingData, topKInfo,
                    isLargest);
            }
        } else {
            if (tmplocalBytes % LOCAL_BYTES != 0) {
                tmplocalBytes = (tmplocalBytes + LOCAL_BYTES - 1) / LOCAL_BYTES * LOCAL_BYTES;
            }
            pipe.InitBuffer(tmplocalBuf, tmplocalBytes);
            AscendC::LocalTensor<uint8_t> tmplocalTensor = tmplocalBuf.Get<uint8_t>();
            if (isSmallMode) {
                AscendC::TopK<T, isInitIndex, isHasfinish, isReuseSrc, AscendC::TopKMode::TOPK_NSMALL>(dstLocalValue,
                    dstLocalIndex, srcLocalValue, srcLocalIndex, srcLocalFinish, tmplocalTensor, 
                    k, topKTilingData, topKInfo, isLargest);
            } else {
                AscendC::TopK<T, isInitIndex, isHasfinish, isReuseSrc, AscendC::TopKMode::TOPK_NORMAL>(dstLocalValue,
                    dstLocalIndex, srcLocalValue, srcLocalIndex, srcLocalFinish, tmplocalTensor, 
                    k, topKTilingData, topKInfo, isLargest);
            }
        }

        outQueueY1.EnQue<T>(dstLocalValue);
        outQueueY2.EnQue<U>(dstLocalIndex);

        inQueueX1.FreeTensor(srcLocalValue);
        inQueueX2.FreeTensor(srcLocalIndex);
        inQueueX3.FreeTensor(srcLocalFinish);
    }
    __aicore__ inline void CopyOut() {
        AscendC::LocalTensor<T> dstLocalValue = outQueueY1.DeQue<T>();
        AscendC::LocalTensor<U> dstLocalIndex = outQueueY2.DeQue<U>();

        AscendC::DataCopy(dstGlobal1, dstLocalValue, outValueDataSize);
        AscendC::DataCopy(dstGlobal2, dstLocalIndex, outIndexDataSize);
        outQueueY1.FreeTensor(dstLocalValue);
        outQueueY2.FreeTensor(dstLocalIndex);
    }

private:
    AscendC::GlobalTensor<T> srcGlobal1;
    AscendC::GlobalTensor<U> srcGlobal2;
    AscendC::GlobalTensor<bool> srcGlobal3;
    AscendC::GlobalTensor<T> dstGlobal1;
    AscendC::GlobalTensor<U> dstGlobal2;

    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECIN, 1> inQueueX1;
    AscendC::TQue<AscendC::TPosition::VECIN, 1> inQueueX2;
    AscendC::TQue<AscendC::TPosition::VECIN, 1> inQueueX3;

    AscendC::TQue<AscendC::TPosition::VECOUT, 1> outQueueY1;
    AscendC::TQue<AscendC::TPosition::VECOUT, 1> outQueueY2;

    AscendC::TBuf<AscendC::TPosition::VECCALC> tmplocalBuf;
    uint32_t tmplocalBytes = 0;
    uint32_t inDataSize = 0;
    uint32_t inputdexDataSize = 0;
    uint32_t inputdexBytes = 0;
    uint32_t finishLocalBytes;
    uint32_t outValueDataSize = 0;
    uint32_t outIndexDataSize = 0;
    uint32_t k;
    uint32_t k_pad;
    uint32_t kpad_index;
    bool isLargest = true;
    TopkTiling topKTilingData;
    uint32_t outter;
    uint32_t inner;
    uint32_t n;
    bool isSmallMode = false;
};

} // namespace MyCustomKernel

#endif // EXAMPLES_SORT_TOPK_CUSTOM_H
